import json

from channels.exceptions import StopConsumer
from channels.generic.websocket import AsyncWebsocketConsumer, WebsocketConsumer
from chatCompletion.entity.serializers.ChatMessageWSDTO import ChatMessageWSDTO
from chatCompletion.service.ChatService import ChatService
from chatCompletion.entity.result.BaseResult import WebSocketResult


class ChatConsumer(AsyncWebsocketConsumer):
    async def websocket_connect(self, message):
        # 有客户端向后端发送websocket连接请求的时候自动自动触发
        # 服务端允许和客户端创建链接
        print(message)
        chat_msg_data = ChatMessageWSDTO(data=message)
        print(chat_msg_data)
        await self.accept()

    async def websocket_receive(self, message):
        print(message)
        # 浏览器基于websocket向后端发送数据，自动触发接收消息
        chat_msg_data: dict = json.loads(ChatMessageWSDTO(data=message.get("text")).initial_data)
        print(chat_msg_data)
        response_data = await self.chat_complete(chat_msg_data)
        await self.send(WebSocketResult.success(data=response_data))
        # self.close()

    async def websocket_disconnect(self, message):
        # 客户端与服务端断开连接，自动触发
        raise StopConsumer()

    async def chat_complete(self, data: dict):
        return await ChatService.simple_chat(data['message'])
